import sys
import numpy as np
from PyQt5.QtWidgets import QApplication, QMainWindow, QVBoxLayout, QWidget, QLabel, QGraphicsView
from PyQt5.QtChart import QChart, QChartView, QLineSeries, QValueAxis
from PyQt5.QtGui import QPen, QColor, QPainter
from PyQt5.QtCore import Qt, QPointF

# 模拟数据
np.random.seed(0)
num_points = 50  # 增加数据点以便图表更明显
well_id = 'A1'
depths = np.linspace(1000, 2000, num_points)
rock_desc = ['砂岩', '泥岩', '灰岩']
analyst = '张三'
date = '2024-06-01'
types = ['a', 'b', 'c']
# 模拟
element_to_simulate = {
    "Mg": (0, 15),
    "Al": (5, 20),
    'Si': (10, 25)
}

data = []

for i in range(num_points):
    row = {
        '井号': well_id,
        '深度': depths[i],
        '岩性简述': np.random.choice(rock_desc),
        '分析人': analyst,
        '分析日期': date,
        '类型': np.random.choice(types),
    }
    for element, (min_val, max_val) in element_to_simulate.items():
        row[element] = np.random.uniform(min_val, max_val)
    data.append(row)


class ChartWidget(QMainWindow):
    def __init__(self, data):
        super().__init__()

        self.data = data
        self.setWindowTitle("井数据可视化 - Qt Charts 地质导向图")
        self.setGeometry(100, 100, 1000, 800)

        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        layout = QVBoxLayout(central_widget)

        self.chart = QChart()
        self.chart.setTitle("地质导向交互折线图")
        self.chart.setAnimationOptions(QChart.SeriesAnimations)

        # 创建坐标轴
        self.axisX = QValueAxis()
        self.axisX.setTitleText("深度")
        self.axisX.setLabelFormat("%.1f")
        self.axisX.setTickCount(10)  # 增加刻度线数量
        min_depth, max_depth = min([d['深度'] for d in data]), max([d['深度'] for d in data])
        self.axisX.setRange(min_depth, max_depth)

        self.axisY = QValueAxis()
        self.axisY.setTitleText("元素含量")
        self.axisY.setLabelFormat("%.1f")
        self.axisY.setTickCount(10)  # 增加刻度线数量
        # 根据模拟数据调整纵坐标范围，考虑所有选定元素的范围
        # elements_to_plot = ['Mg', 'Al', 'Si'] # 默认绘制的元素
        elements_to_plot = list(element_to_simulate.keys())  # 获取模拟字典所有键的列表
        all_values = []
        for elem in elements_to_plot:
            all_values.extend([d[elem] for d in data if elem in d])
        min_y, max_y = min(all_values) if all_values else 0, max(all_values) if all_values else 1
        self.axisY.setRange(min_y - (max_y - min_y) * 0.1, max_y + (max_y - min_y) * 0.1)

        self.chart.addAxis(self.axisX, Qt.AlignBottom)
        self.chart.addAxis(self.axisY, Qt.AlignLeft)

        # 添加数据系列 (折线)
        self.series = {}
        colors = {}
        # 使用所有模拟的元素来分配颜色
        elements_to_draw = list(element_to_simulate.keys())
        for element in elements_to_draw:
            # 生成随机颜色 (RGB)
            r = np.random.randint(0, 256)
            g = np.random.randint(0, 256)
            b = np.random.randint(0, 256)
            colors[element] = QColor(r, g, b)

        # 遍历所有要绘制的元素，创建并添加折线系列
        for element in elements_to_draw:
            # 为当前元素创建一个 QLineSeries 对象，它将存储折线的数据点
            series = QLineSeries()
            # 设置折线的名称，这将显示在图例中 (直接使用元素名称)
            series.setName(element)

            # 获取对应元素的颜色
            # 使用 colors.get() 方法，如果元素在 colors 字典中不存在，则使用默认颜色 QColor(100, 100, 100) (灰色)
            color = colors.get(element, QColor(100, 100, 100))  # 默认灰色
            # 创建一个 QPen 对象，用于设置折线的样式 (颜色和宽度)
            pen = QPen(color)  # 画笔颜色
            # 设置画笔的宽度 (你这里写错了，应该是 setWidth 而不是 setWodth)
            pen.setWidth(2)  # 设置线宽为2像素
            # 将画笔设置给数据系列，确定折线的外观
            series.setPen(pen)

            # 添加数据点 QPointF(x, y) -> QPointF(深度, 元素含量)
            # 创建一个列表推导式，从原始数据中筛选出当前元素的有效数据点 (深度, 含量)
            # 这里的 if element in row 确保数据点字典中确实包含当前元素，增加了健壮性
            valid_points = [(row['深度'], row[element]) for row in self.data if element in row]
            # 遍历有效的点列表，将每个点添加到 QLineSeries 中
            for depth, value in valid_points:
                series.append(QPointF(depth, value))  # QPointF 用于存储点的 (x, y) 坐标

            # 检查是否有有效的数据点，只在有数据点时才将系列添加到图表
            if len(valid_points) > 0:
                # 将创建并填充数据的折线系列添加到 QChart 对象中
                self.chart.addSeries(series)
                # 将横坐标轴 (self.axisX) 关联到这个系列
                series.attachAxis(self.axisX)
                # 将纵坐标轴 (self.axisY) 关联到这个系列
                series.attachAxis(self.axisY)
                # 将创建的系列对象存储在 self.series 字典中，以便后续可以通过元素名称访问 (例如用于更新或交互)
                self.series[element] = series  # 使用元素名称作为 series 字典的键

                # 启用点点击信号，用于触碰显示节点数据
                # 将这个折线系列的 clicked 信号连接到 handlePointClicked 方法
                # 当用户点击这条折线上的某个数据点时，handlePointClicked 方法会被调用，并传入被点击点的坐标
                series.clicked.connect(self.handlePointClicked)  # 使用clicked信号作为触碰事件
                # Note: Qt Charts 没有内置的 hover 事件到具体点，clicked 是一个相对简单的方式演示触碰交互

        # 创建 ChartView
        self.chart_view = QChartView(self.chart)
        self.chart_view.setRenderHint(QPainter.Antialiasing)

        # 设置交互模式：拖动平移
        # self.chart_view.setRubberBand(QChartView.RectangleRubberBand) # 默认矩形缩放
        # setDragMode(QGraphicsView.ScrollHandDrag) 是 QGraphicsView 的方法，QChartView 继承自它
        self.chart_view.setDragMode(QGraphicsView.ScrollHandDrag)  # 启用拖动手势进行平移

        layout.addWidget(self.chart_view)

        # 用于显示节点数据的 QLabel
        self.data_label = QLabel("悬停/点击节点显示数据")
        self.data_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(self.data_label)

    # 处理点点击事件，显示数据
    def handlePointClicked(self, point):
        # 找到是哪个 series 的点被点击了
        clicked_series = self.sender()  # sender() 返回发送信号的对象 (即 QLineSeries)
        if isinstance(clicked_series, QLineSeries):
            # 在 series 中找到被点击的点 (基于坐标)
            # 这是一个简化的查找，对于大量数据可能需要优化
            # 更好的方式是连接 QLineSeries 的 pointsPressed 或 pointsHovered 信号，但 PySide6 的 QLineSeries 没有直接提供简单的 point value 信号
            data_text = f"深度: {point.x():.1f}, 含量: {point.y():.2f}"

            # 更详细的数据查找（需要遍历原始数据）
            matching_rows = [row for row in self.data if abs(row['深度'] - point.x()) < 0.1 and abs(
                row[clicked_series.name().split(' (')[1][:-1]] - point.y()) < 0.1]
            if matching_rows:
                row = matching_rows[0]
                data_text = f"井号: {row['井号']}\n"
                data_text += f"深度: {row['深度']:.1f}\n"
                data_text += f"岩性: {row['岩性简述']}\n"
                data_text += f"类型: {row['类型']}\n"
                # 根据 series name 获取元素，然后显示该元素的值
                element = clicked_series.name().split(' (')[1][:-1]  # 从 "a (Mg)" 提取 "Mg"
                data_text += f"{element}: {row[element]:.2f}\n"
                data_text += f"分析人: {row['分析人']}\n"
                data_text += f"日期: {row['分析日期']}"

            self.data_label.setText(data_text)


def main():
    app = QApplication(sys.argv)
    w = ChartWidget(data)
    w.show()
    sys.exit(app.exec())


if __name__ == '__main__':
    main() 